#! /usr/bin/env python
# -*- coding: utf-8 -*-

import gradio as gr
import numpy as np

from configs import *
from utils.network_utils import get_network
from utils.data_utils import get_dataloader
from train_test import *

pretrained = 'runs/furnace86/furnace86_furnet3/test_exp/checkpoint/train_furnace86_furnet3_test_exp_best.pth.tar'

state = None
# ===== build/load model =====
if pretrained:
    state = torch.load(pretrained)
    model = state['net']
    print('load model finish')
else:
    model = get_network('furnet', 3, 'furnace86')

def predict(i1, i2, i3, i4, i5, i6):
    inp = np.array([i1, i2, i3, i4, i5, i6]) / 100
    inputs = np.array(inp)
    inputs = np.append(inputs, np.mean(inp))
    inputs = np.append(inputs, np.var(inp))
    inputs = torch.Tensor(inputs)
    print(inputs)
    model.eval()
    with torch.no_grad():
        # inputs, targets = inputs.cuda(), targets.cuda()
        outputs = model(inputs)*100
    return outputs.tolist()

inputs = [gr.Slider(0, 1500), gr.Slider(0, 1500), gr.Slider(0, 1500), gr.Slider(0, 1500), gr.Slider(0, 1500), gr.Slider(0, 1500)]
outputs = [gr.Slider(0, 1500), gr.Slider(0, 1500), gr.Slider(0, 1500), gr.Slider(0, 1500), gr.Slider(0, 1500), gr.Slider(0, 1500)]
# outputs = ["number", "number", "number", "number", "number", "number"]

gr.Interface(fn=predict, inputs=inputs, outputs=outputs).launch()

